from evcouplings.couplings import CouplingsModel
import pandas as pd
from evcouplings.mutate import predict_mutation_table, single_mutant_matrix
import numpy as np
import scipy.stats as stats

def calculate_spearman_correlation(X, Y):
    return stats.spearmanr(X, Y)[0]


def calculate_spearman_correlation_p(X, Y):
    return stats.spearmanr(X, Y)[1]

#proteus_names=['UBC9_HUMAN','RASH_HUMAN','TIM_SULSO','P84126_THETH','MTH3_HAEAESTABILIZED','KKA2_KLEPN','BLAT_ECOLX','BG_STRSQ','B3VI55_LIPST','AMIE_PSEAE']

proteus_names=['WW','GFP']
for name in proteus_names:

    #print(name)
    c = CouplingsModel("plmc/example/protein/params/"+name+'.params')
    files=open('proteus/'+name+'.fasta','r')
    seq=files.readlines()
    d=len(seq)
    fs=str()
    for jj in range(1,d):
        t=seq[jj].replace('\n','')
        #print(t)
        #print(seq[jj])
        fs=fs+str(t)
    #print(seq)
    #print(len(seq[1]))
    #print(fs)
    c.target_seq=fs

    data = pd.read_csv('experiment/'+name+'.tsv')
    reals=[]
    zz=np.array(data)
    

    for z in zz:
        z = str(z).replace('[', '').replace(']', '').replace('\'', '').replace('\\', '').split('t')
        
        fenge=z[0].split(';')
        score=float(z[1])
        d_fenge=len(fenge)
        mul=[]
        for qwe in range(d_fenge):
            #print(int(fenge[qwe][1:-2]),fenge[qwe][0],fenge[qwe][-1])
            mul.append((int(fenge[qwe][1:-1]),fenge[qwe][0],fenge[qwe][-1]))
            
        #print(mul)
        delta_E, delta_E_couplings, delta_E_fields = c.delta_hamiltonian(mul)
        
        reals.append([z[0], score,delta_E])
        
        #reals.append([z[0], float(z[1])])
    
    #print(reals)
    data_pred=pd.DataFrame(reals,columns=['mutant','score',"effect_prediction_epistatic"])
    #print(data)
# predict mutations using our model
    
    x=data_pred['score']
    y=data_pred['effect_prediction_epistatic']
    sp_p = calculate_spearman_correlation_p(x, y)
    sp = calculate_spearman_correlation(x, y)
    print(name + ':')
    print(sp, sp_p)
    #print(data_pred.head())
    
    